import cv2
import pywt
import numpy as np
import matplotlib.pyplot as plt


# arnold置换算法，key位置换次数
# 代码实现了Arnold置换算法，该算法是一种用于混淆图像像素位置的方法，用于增强水印的安全性。
# arnold函数执行Arnold置换，deArnold函数执行逆Arnold置换。
def arnold(img):
    r, c = img.shape
    p = np.zeros((r, c), np.uint8)
    a = 1
    b = 1
    for i in range(r):
        for j in range(c):
            x = (i + b * j) % r
            y = (a * i + (a * b + 1) * j) % c
            p[x, y] = img[i, j]
    return p
# 逆arnold置换算法，key位置换次数
def dearnold(img):
    r, c = img.shape
    p = np.zeros((r, c), np.uint8)
    a = 1
    b = 1
    for i in range(r):
        for j in range(c):
            x = ((a * b + 1) * i - b * j) % r
            y = (-a * i + j) % c
            p[x, y] = img[i, j]
    return p

# 计算均方误差函数
# 越小越好
def calculate_mse(image1, image2):
    mse = np.mean((image1 - image2) ** 2)
    return mse
# 计算信噪比函数
# image:嵌入水印的图像
# noise:被攻击后的image
# 越高越好
def calculate_snr(image, noise):
    signal_power = np.mean(image ** 2)
    noise_power = np.mean(noise ** 2)
    snr = 10 * np.log10(signal_power / noise_power)
    return snr
# 噪声攻击 随机噪声大小
def noise_attack(watermarked_image, noise_size):
    # 生成高斯噪声
    noise = np.random.normal(0, noise_size, watermarked_image.shape).astype(np.float32)
    # 添加噪声到水印图像
    noisy_image = watermarked_image + noise
    # 将噪声图像限制在 [0, 255] 范围内
    noisy_image = np.clip(noisy_image, 0, 255).astype(np.uint8)
    return noisy_image

# 裁剪攻击 随机位置 像素罩，不改变图像size
def crop_attack(image, crop_percentage):
    # 获取图像尺寸
    height, width = image.shape[:2]
    # 计算裁剪区域的大小
    crop_height = int(height * crop_percentage)
    crop_width = int(width * crop_percentage)
    # 随机选择裁剪区域的起始点
    x1 = np.random.randint(0, width - crop_width + 1)
    y1 = np.random.randint(0, height - crop_height + 1)
    # 计算裁剪区域的结束点
    x2 = x1 + crop_width
    y2 = y1 + crop_height
    # 执行裁剪攻击
    attacked_image = image.copy()  # 复制原始图像
    attacked_image[y1:y2, x1:x2] = 0  # 将裁剪区域的像素值置零
    return attacked_image
# 旋转攻击
def rotation_attack(image, rotation_angle):
    height, width = image.shape[:2]
    rotation_matrix = cv2.getRotationMatrix2D((width / 2, height / 2), rotation_angle, 1)
    rotated_image = cv2.warpAffine(image, rotation_matrix, (width, height))
    return rotated_image

# 在setwaterMark函数中，首先将载体图像和水印图像调整为所需的大小。
# 然后，对载体图像进行灰度处理，并对水印图像进行Arnold置换。
# 接下来，对载体图像进行三级小波变换，并对水印图像进行一级小波变换。
# 使用自定义的嵌入系数，将水印信息嵌入到小波系数中。最后，对小波系数进行逆小波变换以重构带有水印的图像。
# 返回值是newImg
def setwaterMark(waterTmg, Img, a1, a2, a3, a4, round):
    # print('水印嵌入...')
    Img = cv2.resize(Img, (400, 400))
    waterTmg = cv2.resize(waterTmg, (101, 101))
    # 载体图像灰度处理
    Img1 = cv2.cvtColor(Img, cv2.COLOR_RGB2GRAY)
    waterTmg1 = cv2.cvtColor(waterTmg, cv2.COLOR_RGB2GRAY)
    # 对水印图像进行Arnold变换
    for r in range(round):
        waterTmg1 = arnold(waterTmg1)
    # cv2.imshow('水印+Arnold', Img1)
    # cv2.waitKey(0)
    # 载体图像三级小波变换
    c = pywt.wavedec2(Img1, 'db2', level=3)
    [cl, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)] = c
    # waterTmg1 = cv2.resize(waterTmg1, (101, 101))
    # cv2.imshow('Arnold水印(101,101)', waterTmg1)
    # cv2.waitKey(0)
    # 水印图像一级小波变换
    d = pywt.wavedec2(waterTmg1, 'db2', level=1)
    [ca1, (ch1, cv1, cd1)] = d
    # 嵌入
    cl = cl + ca1 * a1
    cH3 = cH3 + ch1 * a2
    cV3 = cV3 + cv1 * a3
    cD3 = cD3 + cd1 * a4
    # 图像重构
    newImg = pywt.waverec2([cl, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)], 'db2')
    newImg = np.array(newImg, np.uint8)
    mse_embed = calculate_mse(Img1, newImg)
    print('嵌入图像和原图 MSE = ', mse_embed)
    # cv2.imshow("after", newImg)
    # cv2.waitKey(0)
    cv2.imwrite('./after.bmp', newImg)
    return newImg

# 在getwaterMark函数中，对原始图像和带有水印的图像进行相同的小波变换。
# 然后，使用相同的自定义嵌入系数，从小波系数中提取水印信息。
# 最后，执行逆Arnold置换以提取原始水印图像。
# 参数：第一个是原始图像，第二个是嵌入水印后的图像
def getwaterMark(originalImage, Img, a1, a2, a3, a4, round):
    # print('水印提取...')

    originalImage = cv2.resize(originalImage, (400, 400))
    # 原始图像灰度处理
    Img1 = cv2.cvtColor(originalImage, cv2.COLOR_RGB2GRAY)
    Img = cv2.cvtColor(Img, cv2.COLOR_RGB2GRAY)

    # cv2.imshow('ds', Img1)
    # cv2.waitKey(0)

    # 原始图像三级小波变换
    d = pywt.wavedec2(Img1, 'db2', level=3)
    [dl, (dH3, dV3, dD3), (dH2, dV2, dD2), (dH1, dV1, dD1)] = d

    # 载体图像三级小波变换
    c = pywt.wavedec2(Img, 'db2', level=3)
    [cl, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)] = c

    # 嵌入算法逆运算

    ca1 = (cl - dl) / a1
    ch1 = (cH3 - dH3) / a2
    cv1 = (cV3 - dV3) / a3
    cd1 = (cD3 - dD3) / a4

    # 水印图像重构
    waterImg = pywt.waverec2([ca1, (ch1, cv1, cd1)], 'db2')
    waterImg = np.array(waterImg, np.uint8)

    # cv2.imshow("get", waterImg)
    # cv2.waitKey(0)

    # 对提取的水印图像进行逆Arnold变换
    # waterImg = inverse_arnold_transform(waterImg, key)
    for r in range(round):
        waterImg = dearnold(waterImg)

    print('水印提取完成！')
    # waterImg = cv2.resize(waterImg, (400, 400))
    # cv2.imshow("get", get_waterImg)
    # cv2.imwrite('./getWater.bmp', waterImg)
    # cv2.waitKey(0)
    return waterImg

# dct变换
def dct(original_image, watermark, alpha = 0.25):
    dct_image = cv2.dct(np.float32(original_image))    # 执行DCT变换
    dct_watermark = cv2.dct(np.float32(watermark))    # 将水印嵌入到DCT系数中
    watermarked_dct = dct_image + alpha * dct_watermark
    watermarked_image = cv2.idct(watermarked_dct).real    # 执行逆DCT变换
    return watermarked_image
# idct变换
def idct(original_image, watermarked_image, alpha = 0.25):
    watermarked_image_dct = cv2.dct(np.float32(watermarked_image))  # 水印dct
    dct_image = cv2.dct(np.float32(original_image))  # 原图dct
    extracted_dct_watermark = (watermarked_image_dct - dct_image) / alpha
    extracted_watermark = cv2.idct(extracted_dct_watermark).real  # 提取出嵌入的水印
    return extracted_watermark
# dft变换
def dft(original_image, watermark,  alpha = 0.25):
    dft_image = np.fft.fft2(np.float32(original_image))  #原图dft变换
    dft_watermark = np.fft.fft2(np.float32(watermark))  # 水印dft变换
    watermarked_dft = dft_image + alpha * dft_watermark  # 嵌入
    watermarked_image = np.fft.ifft2(watermarked_dft).real  #dft逆变换
    return watermarked_image
# idft变换
def idft(original_image, watermarked_image, alpha = 0.25):
    watermarked_image_dft = np.fft.fft2(np.float32(watermarked_image)) #带水印图片dft
    dft_image = np.fft.fft2(np.float32(original_image)) #原图dft
    extracted_dft_watermark = (watermarked_image_dft - dft_image) / alpha  #水印dft
    extracted_watermark = np.fft.ifft2(extracted_dft_watermark).real # 提取水印逆dft
    return extracted_watermark
# dct+arnold嵌入代码
def set(original_image, watermark, round = 11, alpha = 0.25):
    # 阿尔诺德变换
    for k in range(round):
        watermark = arnold(watermark)
    # watermarked_image = dct(original_image, watermark, alpha)# dct变换
    watermarked_image = dft(original_image, watermark, alpha)# dft变换
    return watermarked_image
# dct+arnold提取代码
def get(original_image, watermarked_image, round = 11, alpha = 0.25):
    # extracted_watermark = idct(original_image, watermarked_image, alpha)# idct变换
    extracted_watermark = idft(original_image, watermarked_image, alpha)# idft变换
    # 阿尔诺德逆变换
    for k in range(round):
        extracted_watermark = dearnold(extracted_watermark)
    return extracted_watermark

if __name__ == '__main__':
    # 读取水印图像、载体图像
    # 并且转化为灰度图像
    originalImage = cv2.imread('yitaji.jpg', cv2.IMREAD_GRAYSCALE)
    waterImg = cv2.imread('lsc_sign.jpg', cv2.IMREAD_GRAYSCALE)
    round = 11
    alpha = 0.25
    # 水印二值化，阈值128
    # _, waterImg = cv2.threshold(waterImg, 128, 255, cv2.THRESH_BINARY)
    # 水印嵌入
    # setwaterMark(waterImg, originalImage, 0.1, 0.2, 0.1, 0.1, 0)

    # 读取原始图像、嵌入水印图像
    # originalImage = cv2.imread('yitaji.jpg')
    # watered_img = cv2.imread('./after.bmp')
    watered_img = set(originalImage, waterImg, round, alpha)

    # 噪声攻击 旋转攻击 裁剪攻击
    noise_img = noise_attack(watered_img, 10)
    rotation_img = rotation_attack(watered_img, 1)
    crop_img = crop_attack(watered_img, 0.5)
    # 图像水印提取 三级小波变换
    # get_waterImg = getwaterMark(originalImage, watered_img, 0.1, 0.2, 0.1, 0.1, 0)
    # get_noise_waterImg = getwaterMark(originalImage, noise_img, 0.1, 0.2, 0.1, 0.1, 2)
    # get_rotation_img = getwaterMark(originalImage, rotation_img, 0.1, 0.2, 0.1, 0.1, 2)
    # get_crop_img = getwaterMark(originalImage, crop_img, 0.1, 0.2, 0.1, 0.1, 2)

    # 简单dct变换 水印提取
    get_waterImg = get(originalImage, watered_img, round, alpha)
    get_noise_waterImg = get(originalImage, noise_img, round, alpha)
    get_rotation_img = get(originalImage, rotation_img, round, alpha)
    get_crop_img = get(originalImage, crop_img, round, alpha)


    # 水印提取和显示
    # cv2.imshow("get_waterImg", get_waterImg)
    # cv2.imshow("get_noise_waterImg", get_noise_waterImg)
    # cv2.imshow("get_rotation_img", get_rotation_img)
    # cv2.imshow("get_crop_img", get_crop_img)
    # cv2.waitKey(0)
    # snr计算 以及打印
    noise_snr = calculate_snr(get_waterImg, get_noise_waterImg)
    rotation_snr = calculate_snr(get_waterImg, get_rotation_img)
    crop_snr = calculate_snr(get_waterImg, get_crop_img)
    print('noise_snr = ', noise_snr)
    print('rotation_snr = ', rotation_snr)
    print('crop_snr = ', crop_snr)

    plt.subplot(4, 3, 1)
    plt.imshow(originalImage, cmap='gray')
    plt.title('originalImage')
    plt.axis('off')

    plt.subplot(4, 3, 2)
    plt.imshow(waterImg, cmap='gray')
    plt.title('waterImg')
    plt.axis('off')

    plt.subplot(4, 3, 3)
    plt.imshow(watered_img, cmap='gray')
    plt.title('watered_img')
    plt.axis('off')

    plt.subplot(4, 3, 4)
    plt.imshow(noise_img, cmap='gray')
    plt.title('noise_img')
    plt.axis('off')

    plt.subplot(4, 3, 5)
    plt.imshow(rotation_img, cmap='gray')
    plt.title('rotation_img')
    plt.axis('off')

    plt.subplot(4, 3, 6)
    plt.imshow(crop_img, cmap='gray')
    plt.title('crop_img')
    plt.axis('off')

    plt.subplot(4, 3, 7)
    plt.imshow(get_waterImg, cmap='gray')
    plt.title('get_waterImg')
    plt.axis('off')

    plt.subplot(4, 3, 8)
    plt.imshow(get_noise_waterImg, cmap='gray')
    plt.title('get_noise_waterImg')
    plt.axis('off')

    plt.subplot(4, 3, 9)
    plt.imshow(get_rotation_img, cmap='gray')
    plt.title('get_rotation_img')
    plt.axis('off')

    plt.subplot(4, 3, 10)
    plt.imshow(get_crop_img, cmap='gray')
    plt.title('get_crop_img')
    plt.axis('off')

    plt.show()

"""---------------
上面的代码是一个用于数字水印嵌入和提取的Python程序，其中使用了小波变换和Arnold置换技术。以下是对代码的详细解释：

1. **Arnold置换算法**：
   - `arnold`函数实现了Arnold置换算法，用于混淆图像像素位置。`img`是输入图像，`key`表示位置置换次数。通过多次迭代，图像的像素位置会被置换，以增强水印的安全性。
   - `deArnold`函数执行逆Arnold置换，用于还原经过Arnold置换的图像。

2. **计算均方误差函数**：
   - `calculate_mse`函数计算两幅图像之间的均方误差（MSE）。MSE用于度量两幅图像之间的差异。

3. **计算信噪比函数**：
   - `calculate_snr`函数计算图像的信噪比（SNR）。SNR用于度量信号与噪声之间的相对强度。

4. **水印嵌入过程**（`setwaterMark`函数）：
   - 加载水印图像和载体图像，将它们调整为所需的大小。
   - 将载体图像转换为灰度图像，并对水印图像执行Arnold置换。
   - 对载体图像进行三级小波变换，对水印图像进行一级小波变换。
   - 使用自定义的嵌入系数，将水印信息嵌入到小波系数中。
   - 对小波系数执行逆小波变换以重构带有水印的图像。
   - 计算并输出水印嵌入后的图像的均方误差（MSE）。

5. **水印提取过程**（`getwaterMark`函数）：
   - 加载原始图像和嵌入了水印的图像。
   - 对这两幅图像分别进行相同的小波变换。
   - 使用相同的自定义嵌入系数，从小波系数中提取水印信息。
   - 对提取的小波系数执行逆小波变换以还原水印图像。
   - 对还原的水印图像执行逆Arnold置换以还原原始水印图像。
   - 计算并输出水印提取后的图像的均方误差（MSE）。

6. **主程序**：
   - 主程序首先加载水印图像和载体图像。
   - 调用`setwaterMark`函数进行水印嵌入，并将嵌入后的图像保存为`after.bmp`。
   - 再次加载原始图像和嵌入了水印的图像。
   - 调用`getwaterMark`函数进行水印提取，并将提取的水印图像显示出来。
    此代码结合了Arnold置换和小波变换，用于实现数字水印的嵌入和提取。水印嵌入和提取过程分别使用了小波变换，以保护水印并提供鲁棒性。
    均方误差（MSE）用于评估水印的嵌入和提取效果，信噪比（SNR）用于评估图像质量和鲁棒性。
----------------"""
